#include "yolov5_dnn.h"

int main() {
    // 加载 ONNX 模型
    std::string modelPath = "/home/lfj/tmp/opencv/c++/models/best.onnx";  // 替换为你的模型路径
    cv::dnn::Net net = cv::dnn::readNetFromONNX(modelPath);
    
    if (net.empty()) {
        std::cerr << "Failed to load the ONNX model!" << std::endl;
        return -1;
    }else{
        std::cout << "Success to load the ONNX model!" << std::endl;
    }

    // 加载输入图像
    std::string imagePath = "/home/lfj/tmp/opencv/c++/data/val.png";  // 替换为你的输入图片路径
    cv::Mat image = cv::imread(imagePath, cv::IMREAD_COLOR);
    if (image.empty()) {
        std::cerr << "Failed to load the image!" << std::endl;
        return -1;
    }else{
        std::cout << "Success to load the image!" << std::endl;
    }

    int col = image.cols;
    int row = image.rows;
    int _max = MAX(col, row);
    cv::Mat resized = cv::Mat::zeros(_max, _max, CV_8UC3);
    image.copyTo(resized(cv::Rect(0, 0, col, row)));

    float x_factor = resized.cols / 640.0f;
    float y_factor = resized.rows / 640.0f;

    // 预处理图像
    cv::Mat blob;
    cv::dnn::blobFromImage(image, blob, 1.0 / 255.0, cv::Size(640, 640), cv::Scalar(0, 0, 0), true, false); // 根据模型要求调整参数
    net.setInput(blob);

#if 0
    cv::Mat preds = net.forward();
 
    // 后处理, 1x25200x85
    // std::cout << "rows: "<< preds.size[1]<< " data: " << preds.size[2] << std::endl;
    cv::Mat det_output(preds.size[1], preds.size[2], CV_32F, preds.ptr<float>());
    std::vector<cv::Rect> boxes;
    std::vector<int> classIds;
    std::vector<float> confidences;
    for (int i = 0; i < det_output.rows; i++) {
        float confidence = det_output.at<float>(i, 4);
        if (confidence < 0.5) {
            continue;
        }
        cv::Mat classes_scores = det_output.row(i).colRange(5, 85);
        cv::Point classIdPoint;
        double score;
        minMaxLoc(classes_scores, 0, &score, 0, &classIdPoint);
 
        // 置信度 0～1之间
        if (score > this->threshold_score)
        {
            float cx = det_output.at<float>(i, 0);
            float cy = det_output.at<float>(i, 1);
            float ow = det_output.at<float>(i, 2);
            float oh = det_output.at<float>(i, 3);
            int x = static_cast<int>((cx - 0.5 * ow) * x_factor);
            int y = static_cast<int>((cy - 0.5 * oh) * y_factor);
            int width = static_cast<int>(ow * x_factor);
            int height = static_cast<int>(oh * y_factor);
            cv::Rect box;
            box.x = x;
            box.y = y;
            box.width = width;
            box.height = height;
 
            boxes.push_back(box);
            classIds.push_back(classIdPoint.x);
            confidences.push_back(score);
        }
    }
 
    // NMS
    std::vector<int> indexes;
    cv::dnn::NMSBoxes(boxes, confidences, 0.25, 0.45, indexes);
    for (size_t i = 0; i < indexes.size(); i++) {
        DetectResult dr;
        int index = indexes[i];
        int idx = classIds[index];
        dr.box = boxes[index];
        dr.classId = idx;
        dr.score = confidences[index];
        cv::rectangle(frame, boxes[index], cv::Scalar(0, 0, 255), 2, 8);
        cv::rectangle(frame, cv::Point(boxes[index].tl().x, boxes[index].tl().y - 20), 
            cv::Point(boxes[index].br().x, boxes[index].tl().y), cv::Scalar(0, 255, 255), -1);
        results.push_back(dr);
    }
 
    std::ostringstream ss;
    std::vector<double> layersTimings;
    double freq = cv::getTickFrequency() / 1000.0;
    double time = net.getPerfProfile(layersTimings) / freq;
    ss << "FPS: " << 1000 / time << " ; time : " << time << " ms";
    putText(frame, ss.str(), cv::Point(20, 40), cv::FONT_HERSHEY_PLAIN, 2.0, cv::Scalar(255, 0, 0), 2, 8);
#else
    // 进行推理
    std::vector<cv::Mat> output;
    net.forward(output);

    std::cout << "Number of outputs: " << output.size() << std::endl;
    for (size_t i = 0; i < output.size(); ++i) {
        std::cout << "Output " << i << " shape: " << output[i].size << std::endl;
    }

    // 处理输出（根据模型的输出结构进行处理）
    // 假设输出是一个边界框（这是一个示例，根据你的模型进行调整）
    for (const auto& mat : output) {
        std::cout << "mat.rows = " << mat.rows << std::endl;
        for (int i = 0; i < mat.rows; i++) {
            float confidence = mat.at<float>(i, 4); // 假设第5列是置信度
            std::cout << "Confidence: " << confidence << std::endl;
            if (confidence > 0.5) { // 置信度阈值
                int left = static_cast<int>(mat.at<float>(i, 0) * image.cols);
                int top = static_cast<int>(mat.at<float>(i, 1) * image.rows);
                int right = static_cast<int>(mat.at<float>(i, 2) * image.cols);
                int bottom = static_cast<int>(mat.at<float>(i, 3) * image.rows);
                
                // 绘制边界框
                cv::rectangle(image, cv::Point(left, top), cv::Point(right, bottom), cv::Scalar(0, 255, 0), 2);
                std::cout << "Detected object with confidence: " << confidence << std::endl;
            }
        }
    }
#endif

    // 保存结果图像
    std::string outputImagePath = "output.png"; // 输出图像路径
    cv::imwrite(outputImagePath, image);
    std::cout << "Output image saved as: " << outputImagePath << std::endl;

    return 0;
}